/* $Id: VBoxMPRegistry.cpp $ */
/** @file
 * VBox WDDM Miniport registry related functions
 */

/*
 * Copyright (C) 2011-2020 Oracle Corporation
 *
 * This file is part of VirtualBox Open Source Edition (OSE), as
 * available from http://www.virtualbox.org. This file is free software;
 * you can redistribute it and/or modify it under the terms of the GNU
 * General Public License (GPL) as published by the Free Software
 * Foundation, in version 2 as it comes in the "COPYING" file of the
 * VirtualBox OSE distribution. VirtualBox OSE is distributed in the
 * hope that it will be useful, but WITHOUT ANY WARRANTY of any kind.
 */

#include <WinError.h>
#include "VBoxMPWddm.h"
#include "VBoxMPRegistry.h"

NTSTATUS vboxWddmRegQueryDrvKeyName(PVBOXMP_DEVEXT pDevExt, ULONG cbBuf, PWCHAR pBuf, PULONG pcbResult)
{
    WCHAR fallBackBuf[2];
    PWCHAR pSuffix;
    BOOLEAN bFallback = false;

    if (cbBuf > sizeof(VBOXWDDM_REG_DRVKEY_PREFIX))
    {
        memcpy(pBuf, VBOXWDDM_REG_DRVKEY_PREFIX, sizeof (VBOXWDDM_REG_DRVKEY_PREFIX));
        pSuffix = pBuf + (sizeof (VBOXWDDM_REG_DRVKEY_PREFIX) - 2) / 2;
        cbBuf -= sizeof (VBOXWDDM_REG_DRVKEY_PREFIX) - 2;
    }
    else
    {
        pSuffix = fallBackBuf;
        cbBuf = sizeof (fallBackBuf);
        bFallback = true;
    }

    NTSTATUS Status = IoGetDeviceProperty(pDevExt->pPDO,
                        DevicePropertyDriverKeyName,
                        cbBuf,
                        pSuffix,
                        &cbBuf);
    if (Status == STATUS_SUCCESS && bFallback)
        Status = STATUS_BUFFER_TOO_SMALL;
    if (Status == STATUS_BUFFER_TOO_SMALL)
        *pcbResult = cbBuf + sizeof (VBOXWDDM_REG_DRVKEY_PREFIX) - 2;

    return Status;
}

NTSTATUS vboxWddmRegOpenKeyEx(OUT PHANDLE phKey, IN HANDLE hRootKey, IN PWCHAR pName, IN ACCESS_MASK fAccess)
{
    OBJECT_ATTRIBUTES ObjAttr;
    UNICODE_STRING RtlStr;

    RtlInitUnicodeString(&RtlStr, pName);
    InitializeObjectAttributes(&ObjAttr, &RtlStr, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, hRootKey, NULL);

    return ZwOpenKey(phKey, fAccess, &ObjAttr);
}

NTSTATUS vboxWddmRegCreateKeyEx(OUT PHANDLE phKey, IN HANDLE hRootKey, IN PWCHAR pName, IN ACCESS_MASK fAccess)
{
    OBJECT_ATTRIBUTES ObjAttr;
    UNICODE_STRING RtlStr;
    ULONG Disposition;

    RtlInitUnicodeString(&RtlStr, pName);
    InitializeObjectAttributes(&ObjAttr, &RtlStr, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, hRootKey, NULL);

    return ZwCreateKey(phKey, fAccess, &ObjAttr, 0, &RtlStr, REG_OPTION_NON_VOLATILE, &Disposition);
}


NTSTATUS vboxWddmRegCreateKey(OUT PHANDLE phKey, IN PWCHAR pName, IN ACCESS_MASK fAccess)
{
    return vboxWddmRegCreateKeyEx(phKey, NULL, pName, fAccess);
}

NTSTATUS vboxWddmRegOpenKey(OUT PHANDLE phKey, IN PWCHAR pName, IN ACCESS_MASK fAccess)
{
    return vboxWddmRegOpenKeyEx(phKey, NULL, pName, fAccess);
}

NTSTATUS vboxWddmRegQueryValueDword(IN HANDLE hKey, IN PWCHAR pName, OUT PDWORD pDword)
{
    struct
    {
        KEY_VALUE_PARTIAL_INFORMATION Info;
        UCHAR Buf[32]; /* should be enough */
    } Buf;
    ULONG cbBuf;
    UNICODE_STRING RtlStr;
    RtlInitUnicodeString(&RtlStr, pName);
    NTSTATUS Status = ZwQueryValueKey(hKey,
                &RtlStr,
                KeyValuePartialInformation,
                &Buf.Info,
                sizeof(Buf),
                &cbBuf);
    if (Status == STATUS_SUCCESS)
    {
        if (Buf.Info.Type == REG_DWORD)
        {
            *pDword = *((PULONG)Buf.Info.Data);
            return STATUS_SUCCESS;
        }
    }

    return STATUS_INVALID_PARAMETER;
}

NTSTATUS vboxWddmRegSetValueDword(IN HANDLE hKey, IN PWCHAR pName, IN DWORD val)
{
    UNICODE_STRING RtlStr;
    RtlInitUnicodeString(&RtlStr, pName);
    return ZwSetValueKey(hKey, &RtlStr,
            NULL, /* IN ULONG  TitleIndex  OPTIONAL, reserved */
            REG_DWORD,
            &val,
            sizeof(val));
}

VP_STATUS VBoxMPCmnRegInit(IN PVBOXMP_DEVEXT pExt, OUT VBOXMPCMNREGISTRY *pReg)
{
    WCHAR Buf[512];
    ULONG cbBuf = sizeof(Buf);
    NTSTATUS Status = vboxWddmRegQueryDrvKeyName(pExt, cbBuf, Buf, &cbBuf);
    if (Status == STATUS_SUCCESS)
    {
        Status = vboxWddmRegOpenKey(pReg, Buf, GENERIC_READ | GENERIC_WRITE);
        if(Status == STATUS_SUCCESS)
            return NO_ERROR;
    }

    /* fall-back to make the subsequent VBoxVideoCmnRegXxx calls treat the fail accordingly
     * basically needed to make as less modifications to the current XPDM code as possible */
    *pReg = NULL;

    return ERROR_INVALID_PARAMETER;
}

VP_STATUS VBoxMPCmnRegFini(IN VBOXMPCMNREGISTRY Reg)
{
    if (!Reg)
    {
        return ERROR_INVALID_PARAMETER;
    }

    NTSTATUS Status = ZwClose(Reg);
    return Status == STATUS_SUCCESS ? NO_ERROR : ERROR_INVALID_PARAMETER;
}

VP_STATUS VBoxMPCmnRegQueryDword(IN VBOXMPCMNREGISTRY Reg, PWSTR pName, UINT32 *pVal)
{
    /* seems like the new code assumes the Reg functions zeroes up the value on failure */
    *pVal = 0;

    if (!Reg)
    {
        return ERROR_INVALID_PARAMETER;
    }

    NTSTATUS Status = vboxWddmRegQueryValueDword(Reg, pName, (PDWORD)pVal);
    return Status == STATUS_SUCCESS ? NO_ERROR : ERROR_INVALID_PARAMETER;
}

VP_STATUS VBoxMPCmnRegSetDword(IN VBOXMPCMNREGISTRY Reg, PWSTR pName, UINT32 Val)
{
    if (!Reg)
    {
        return ERROR_INVALID_PARAMETER;
    }

    NTSTATUS Status = vboxWddmRegSetValueDword(Reg, pName, Val);
    return Status == STATUS_SUCCESS ? NO_ERROR : ERROR_INVALID_PARAMETER;
}

NTSTATUS vboxWddmRegQueryVideoGuidString(PVBOXMP_DEVEXT pDevExt, ULONG cbBuf, PWCHAR pBuf, PULONG pcbResult)
{
    BOOLEAN fNewMethodSucceeded = FALSE;
    HANDLE hKey = NULL;
    NTSTATUS Status = IoOpenDeviceRegistryKey(pDevExt->pPDO, PLUGPLAY_REGKEY_DEVICE, GENERIC_READ, &hKey);
    if (NT_SUCCESS(Status))
    {
        struct
        {
            KEY_VALUE_PARTIAL_INFORMATION Info;
            UCHAR Buf[1024]; /* should be enough */
        } KeyData;
        ULONG cbResult;
        UNICODE_STRING RtlStr;
        RtlInitUnicodeString(&RtlStr, L"VideoID");
        Status = ZwQueryValueKey(hKey,
                    &RtlStr,
                    KeyValuePartialInformation,
                    &KeyData.Info,
                    sizeof(KeyData),
                    &cbResult);
        if (NT_SUCCESS(Status))
        {
            if (KeyData.Info.Type == REG_SZ)
            {
                fNewMethodSucceeded = TRUE;
                *pcbResult = KeyData.Info.DataLength + 2;
                if (cbBuf >= KeyData.Info.DataLength)
                {
                    memcpy(pBuf, KeyData.Info.Data, KeyData.Info.DataLength + 2);
                    Status = STATUS_SUCCESS;
                }
                else
                    Status = STATUS_BUFFER_TOO_SMALL;
            }
        }
        else
        {
            LOGW("ZwQueryValueKey failed, Status 0x%x\n", Status);
        }

        ZwClose(hKey);
    }
    else
    {
        LOGW("IoOpenDeviceRegistryKey failed Status 0x%x\n", Status);
    }

    if (fNewMethodSucceeded)
        return Status;
    else
        LOGW("failed to acquire the VideoID, falling back to the old impl\n");

    Status = vboxWddmRegOpenKey(&hKey, VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY, GENERIC_READ);
    if (Status != STATUS_SUCCESS)
    {
        return Status;
    }
    struct
    {
        KEY_BASIC_INFORMATION Name;
        WCHAR Buf[256];
    } Buf;
    WCHAR KeyBuf[sizeof (VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY)/2 + 256 + 64];
    wcscpy(KeyBuf, VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY);
    ULONG ResultLength;
    BOOL bFound = FALSE;
    for (ULONG i = 0; !bFound; ++i)
    {
        RtlZeroMemory(&Buf, sizeof (Buf));
        Status = ZwEnumerateKey(hKey, i, KeyBasicInformation, &Buf, sizeof (Buf), &ResultLength);
        /* we should not encounter STATUS_NO_MORE_ENTRIES here since this would mean we did not find our entry */
        if (Status != STATUS_SUCCESS)
            break;

        HANDLE hSubKey;
        PWCHAR pSubBuf = KeyBuf + (sizeof (VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY) - 2)/2;
        memcpy(pSubBuf, Buf.Name.Name, Buf.Name.NameLength);
        pSubBuf += Buf.Name.NameLength/2;
        memcpy(pSubBuf, VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY_SUBKEY, sizeof (VBOXWDDM_REG_DISPLAYSETTINGSVIDEOKEY_SUBKEY));
        Status = vboxWddmRegOpenKey(&hSubKey, KeyBuf, GENERIC_READ);
        if (Status != STATUS_SUCCESS)
        {
            break;
        }
        struct
        {
            KEY_VALUE_PARTIAL_INFORMATION Info;
            UCHAR Buf[sizeof (VBOX_WDDM_DRIVERNAME)]; /* should be enough */
        } KeyData;
        ULONG cbResult;
        UNICODE_STRING RtlStr;
        RtlInitUnicodeString(&RtlStr, L"Service");
        Status = ZwQueryValueKey(hSubKey,
                    &RtlStr,
                    KeyValuePartialInformation,
                    &KeyData.Info,
                    sizeof(KeyData),
                    &cbResult);

        if (Status != STATUS_SUCCESS)
        {
            ZwClose(hSubKey);
            ZwClose(hKey);
            return Status;
        }

        if (KeyData.Info.Type != REG_SZ ||
            KeyData.Info.DataLength != sizeof (VBOX_WDDM_DRIVERNAME))
        {
            ZwClose(hSubKey);
            ZwClose(hKey);
            return Status;
        }

        if (!wcscmp(VBOX_WDDM_DRIVERNAME, (PWCHAR)KeyData.Info.Data))
        {
            bFound = TRUE;
            *pcbResult = Buf.Name.NameLength + 2;
            if (cbBuf >= Buf.Name.NameLength + 2)
            {
                memcpy(pBuf, Buf.Name.Name, Buf.Name.NameLength + 2);
            }
            else
            {
                Status = STATUS_BUFFER_TOO_SMALL;
            }
        }

        ZwClose(hSubKey);
    }
    ZwClose(hKey);

    return Status;
}

UNICODE_STRING* vboxWddmVGuidGet(PVBOXMP_DEVEXT pDevExt)
{
    if (pDevExt->VideoGuid.Buffer)
        return &pDevExt->VideoGuid;

    WCHAR VideoGuidBuf[512];
    ULONG cbVideoGuidBuf = sizeof (VideoGuidBuf);
    NTSTATUS Status = vboxWddmRegQueryVideoGuidString(pDevExt ,cbVideoGuidBuf, VideoGuidBuf, &cbVideoGuidBuf);
    if (Status == STATUS_SUCCESS)
    {
        PWCHAR pBuf = (PWCHAR)vboxWddmMemAllocZero(cbVideoGuidBuf);
        if (pBuf)
        {
            memcpy(pBuf, VideoGuidBuf, cbVideoGuidBuf);
            RtlInitUnicodeString(&pDevExt->VideoGuid, pBuf);
            return &pDevExt->VideoGuid;
        }
    }

    return NULL;
}

VOID vboxWddmVGuidFree(PVBOXMP_DEVEXT pDevExt)
{
    if (pDevExt->VideoGuid.Buffer)
    {
        vboxWddmMemFree(pDevExt->VideoGuid.Buffer);
        pDevExt->VideoGuid.Buffer = NULL;
    }
}